{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IafxybMjKfBO"
   },
   "source": [
    "##### Copyright 2020 Google"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "pc1aHcGvKmHe"
   },
   "outputs": [],
   "source": [
    "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0bcc53bffae2"
   },
   "source": [
    "# Making Molecule Files for HFVQE\n",
    "\n",
    "In this tutorial we describe the process of making the molecular data files necessary to run the HFVQE code.  We focus on how to use the OpenFermion plugin modules to generate molecular files with canonical Hartree-Fock and generate integrals in a given atomic orbital basis set.  We also provide helper functions to run variational Hartree-Fock simulating the experiment and generating initial parameters.\n",
    "\n",
    "This tutorial will follow the code in `recirq/hfvqe/molecular_data/` for constructing `MolecularData` objects and getting atomic orbital integrals.\n",
    "\n",
    "In addition to the standard requirement of ReCirq and its dependencies, this notebook uses OpenFermion-pyscf (and pyscf) to compute some quantities. We install it below if you don't already have it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FQEYY3gnK51d"
   },
   "source": [
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://www.example.org/cirq/experiments/hfvqe/molecular_data\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on QuantumLib</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/quantumlib/ReCirq/blob/master/docs/hfvqe/molecular_data.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://github.com/quantumlib/ReCirq/blob/master/docs/hfvqe/molecular_data.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://storage.googleapis.com/tensorflow_docs/ReCirq/docs/hfvqe/molecular_data.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8d4ac441b1fc"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "10a67cf51bbe"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    import recirq\n",
    "except ImportError:\n",
    "    !pip install --quiet git+https://github.com/quantumlib/ReCirq\n",
    "        \n",
    "try:\n",
    "    import openfermionpyscf\n",
    "except ImportError:\n",
    "    !pip install --quiet openfermionpyscf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "223e437e1cba"
   },
   "source": [
    "Now we can import the packages required for this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "a78e0799b80a"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import scipy\n",
    "\n",
    "from recirq.hfvqe.molecular_data.molecular_data_construction import (\n",
    "    h6_linear_molecule, h8_linear_molecule, \n",
    "    h10_linear_molecule, h12_linear_molecule, \n",
    "    get_ao_integrals)\n",
    "from recirq.hfvqe.gradient_hf import rhf_minimization, rhf_func_generator\n",
    "from recirq.hfvqe.objective import RestrictedHartreeFockObjective, generate_hamiltonian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3ea4258dbde5"
   },
   "outputs": [],
   "source": [
    "def make_rhf_objective(molecule):\n",
    "    S, Hcore, TEI = get_ao_integrals(molecule)\n",
    "    _, X = scipy.linalg.eigh(Hcore, S)\n",
    "\n",
    "    molecular_hamiltonian = generate_hamiltonian(\n",
    "        general_basis_change(Hcore, X, (1, 0)),\n",
    "        numpy.einsum('psqr', general_basis_change(TEI, X, (1, 0, 1, 0)),\n",
    "                     molecule.nuclear_repulsion))\n",
    "\n",
    "    rhf_objective = RestrictedHartreeFockObjective(molecular_hamiltonian,\n",
    "                                                   molecule.n_electrons)\n",
    "    return rhf_objective, S, Hcore, TEI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "442f67808f6d"
   },
   "source": [
    "## Hydrogen Chain MolecularData\n",
    "\n",
    "We provide helper functions in the `hfvqe` module to generate the Hydrogen chain data.  Each chain is constructed using OpenFermion and Psi4 via the OpenFermion-Psi4 plugin.  We will use creating H6 with equal spacing between Hydrogen atoms as an example.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "17b8ac650bc5"
   },
   "outputs": [],
   "source": [
    "from openfermion import MolecularData, general_basis_change\n",
    "from openfermionpyscf import run_pyscf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "d035f38d8522"
   },
   "outputs": [],
   "source": [
    "n_hydrogens = 6\n",
    "bond_distance = 1.3 # in Angstroms\n",
    "molecule = MolecularData(\n",
    "    geometry=[('H', (0, 0, i * bond_distance)) for i in range(n_hydrogens)],\n",
    "    charge=0,\n",
    "    basis='6-31g',\n",
    "    multiplicity=1,\n",
    "    description=f\"linear_r-{bond_distance}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4eb1c6753dd8"
   },
   "source": [
    "The previous lines set up the MolecularData file.  We can now use pyscf to either run a full self-consistent-field Hartree-Fock calculation or get atomic integrals.  Via Openfermion-Pyscf we provide an interface to running Hartree-Fock, coupled-cluster, second order perturbation theory, configuration-interaction singles-doubles, and full configuration interaction.  Many of these methods depend on parameters such as convergence criteria or initial vectors in the subspace expansion.  `run_pyscf` assumes common defaults which are appropriate for most systems. Below we will run default Hartree-Fock and CISD."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "44f2f2c32fc1"
   },
   "outputs": [],
   "source": [
    "molecule = run_pyscf(molecule, run_scf=True, run_cisd=True)\n",
    "print(\"Hartree-Fock energy:\", molecule.hf_energy, \n",
    "      \"\\nCISD energy:\", molecule.cisd_energy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "880b63e82149"
   },
   "source": [
    "The `MolecularData` file holds almost all information that is required for post-Hartree-Fock correlated calculations.  For example, we provide access to integrals as attributes of `MolecularData`.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0946f58621d7"
   },
   "outputs": [],
   "source": [
    "print(\"Overlap Ints\")\n",
    "print(molecule.overlap_integrals)\n",
    "\n",
    "print()\n",
    "print(\"One-electron integrals\")\n",
    "print(molecule.one_body_integrals)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4f4084b87f8d"
   },
   "source": [
    "For the Hartree-Fock experiment we will need to get the atomic basis integrals from the molecular integrals.  We can use the identity $C^{\\dagger}SC = I$ to reverse the transformation on the one and two electron integrals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "07231b99d404"
   },
   "outputs": [],
   "source": [
    "oei_mo,  tei_mo = molecule.one_body_integrals, molecule.two_body_integrals\n",
    "C = molecule.canonical_orbitals\n",
    "S = molecule.overlap_integrals\n",
    "oei_ao = general_basis_change(oei_mo, C.conj().T @ S, key=(1, 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5d07399e03ef"
   },
   "outputs": [],
   "source": [
    "print(oei_ao)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b18c243f3b5f"
   },
   "outputs": [],
   "source": [
    "print(oei_mo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "845f9a984c27"
   },
   "outputs": [],
   "source": [
    "# Use pyscf to get atomic integrals and compare to transformed integrals from above\n",
    "pyscf_mol = molecule._pyscf_data['mol']\n",
    "t = pyscf_mol.intor('int1e_kin')\n",
    "v = pyscf_mol.intor('int1e_nuc')\n",
    "hcore = t + v\n",
    "assert np.allclose(hcore, oei_ao)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8de5ca003051"
   },
   "source": [
    " The two-electron integrals can also be acquired from the `pyscf_mol` object or we can use the `general_basis_change` to transform the two-electron integrals back into the AO basis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b4766838c0b8"
   },
   "outputs": [],
   "source": [
    "# Transform the two-electron integrals\n",
    "tei_ao = general_basis_change(tei_mo, C.conj().T @ S, key=(1, 1, 0, 0))  \n",
    "\n",
    "# re-ordering for chem->physics storage of the integrals\n",
    "eri_ao = np.einsum('ijkl->iklj', pyscf_mol.intor('int2e', aosym='s1'))  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "e999ce76ba74"
   },
   "outputs": [],
   "source": [
    "assert np.allclose(tei_ao, eri_ao)\n",
    "assert not np.allclose(tei_ao, tei_mo)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0e8dd3aa5d98"
   },
   "source": [
    "We also provide a function in `recirq.hfvqe.molecular_data.molecular_data_construction` that uses psi4 to generate atomic integrals for the HF-VQE study.  Once the atomic orbital integrals are obtained we can perform the first step in setting up the HF-VQE experiment.  This involes transforming the integrals to the core-orbital basis and building an `RestrictedHartreeFockObjective`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cf87a05f418b"
   },
   "outputs": [],
   "source": [
    "# diagonalize the AO-core Hamiltonian\n",
    "_, X = scipy.linalg.eigh(oei_ao, S)  \n",
    "\n",
    "obi = general_basis_change(oei_ao, X, (1, 0))\n",
    "tbi = np.einsum('psqr', general_basis_change(pyscf_mol.intor('int2e', aosym='s1'), X, (1, 0, 1, 0)))\n",
    "molecular_hamiltonian = generate_hamiltonian(obi, tbi,\n",
    "                                             molecule.nuclear_repulsion)\n",
    "\n",
    "rhf_objective = RestrictedHartreeFockObjective(molecular_hamiltonian,\n",
    "                                               molecule.n_electrons)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fd15f30cb136"
   },
   "source": [
    "To get initial parameters we can simulate the Hartree-Fock experiment which is performing variational-Hartree-Fock theory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6347a656b004"
   },
   "outputs": [],
   "source": [
    "from recirq.hfvqe.gradient_hf import rhf_minimization, rhf_func_generator\n",
    "# Uses conjugate gradient to solve non-linear hartree-fock functional:\n",
    "scipy_result = rhf_minimization(rhf_objective, verbose=True)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cf9e3314bae3"
   },
   "outputs": [],
   "source": [
    "print(molecule.hf_energy, scipy_result.fun)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "94bec5ae7352"
   },
   "outputs": [],
   "source": [
    "print(\"Initial Parameters for HF-VQE Study \", scipy_result.x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cd3195881050"
   },
   "outputs": [],
   "source": [
    "# Explicity build kappa matrix\n",
    "from recirq.hfvqe.circuits import rhf_params_to_matrix\n",
    "import matplotlib.pyplot as plt\n",
    "kappa = rhf_params_to_matrix(scipy_result.x, len(rhf_objective.occ) + len(rhf_objective.virt), rhf_objective.occ,\n",
    "            rhf_objective.virt)\n",
    "plt.imshow(kappa)\n",
    "plt.colorbar()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "molecular_data.ipynb",
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
